package data

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/Shopify/sarama"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// DataProcessor 结构体处理数据的流转，从接收MQTT数据到发送到Kafka再到存储到MongoDB
type DataProcessor struct {
	kafkaProducer sarama.SyncProducer
	mongoClient   *mongo.Client
	dbName        string
	dbCollection  string
}

// NewDataProcessor 创建一个新的数据处理器
func NewDataProcessor(kafkaBrokers []string, mongoURI string, dbName string, collection string) (*DataProcessor, error) {
	// 初始化Kafka生产者
	kafkaConfig := sarama.NewConfig()
	kafkaConfig.Producer.Retry.Max = 5
	kafkaConfig.Producer.RequiredAcks = sarama.WaitForAll
	kafkaConfig.Producer.Return.Successes = true

	producer, err := sarama.NewSyncProducer(kafkaBrokers, kafkaConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create kafka producer: %v", err)
	}

	// 初始化MongoDB客户端
	clientOptions := options.Client().ApplyURI(mongoURI)
	client, err := mongo.Connect(context.Background(), clientOptions)
	if err != nil {
		producer.Close()
		return nil, fmt.Errorf("failed to connect to mongodb: %v", err)
	}

	// 验证MongoDB连接
	err = client.Ping(context.Background(), nil)
	if err != nil {
		client.Disconnect(context.Background())
		producer.Close()
		return nil, fmt.Errorf("failed to ping mongodb: %v", err)
	}

	return &DataProcessor{
		kafkaProducer: producer,
		mongoClient:   client,
		dbName:        dbName,
		dbCollection:  collection,
	}, nil
}

// ProcessData 处理从MQTT接收到的数据
func (dp *DataProcessor) ProcessData(deviceID string, topic string, data map[string]interface{}) error {
	// 增强数据，添加设备ID、时间戳等信息
	enhancedData := map[string]interface{}{
		"device_id":   deviceID,
		"topic":       topic,
		"timestamp":   time.Now().Unix(),
		"received_at": time.Now(),
	}

	// 合并原始数据
	for k, v := range data {
		enhancedData[k] = v
	}

	// 将数据转换为JSON
	jsonData, err := json.Marshal(enhancedData)
	if err != nil {
		return fmt.Errorf("failed to marshal data: %v", err)
	}

	// 发送数据到Kafka
	err = dp.sendToKafka("iot-device-data", jsonData)
	if err != nil {
		log.Printf("Warning: failed to send data to kafka: %v", err)
		// 即使Kafka发送失败，也继续尝试存储到MongoDB
	}

	// 存储数据到MongoDB
	err = dp.storeToMongoDB(enhancedData)
	if err != nil {
		return fmt.Errorf("failed to store data to mongodb: %v", err)
	}

	log.Printf("Processed data from device %s successfully", deviceID)
	return nil
}

// sendToKafka 发送数据到Kafka
func (dp *DataProcessor) sendToKafka(topic string, data []byte) error {
	msg := &sarama.ProducerMessage{
		Topic: topic,
		Value: sarama.StringEncoder(data),
	}

	partition, offset, err := dp.kafkaProducer.SendMessage(msg)
	if err != nil {
		return fmt.Errorf("failed to send message to kafka: %v", err)
	}

	log.Printf("Message sent to kafka: topic=%s, partition=%d, offset=%d", topic, partition, offset)
	return nil
}

// storeToMongoDB 存储数据到MongoDB
func (dp *DataProcessor) storeToMongoDB(data map[string]interface{}) error {
	collection := dp.mongoClient.Database(dp.dbName).Collection(dp.dbCollection)

	// 插入文档
	_, err := collection.InsertOne(context.Background(), data)
	if err != nil {
		return fmt.Errorf("failed to insert document into mongodb: %v", err)
	}

	log.Printf("Data stored to mongodb successfully")
	return nil
}

// Close 关闭资源
func (dp *DataProcessor) Close() error {
	// 关闭Kafka生产者
	err := dp.kafkaProducer.Close()
	if err != nil {
		log.Printf("Warning: failed to close kafka producer: %v", err)
	}

	// 关闭MongoDB客户端
	err = dp.mongoClient.Disconnect(context.Background())
	if err != nil {
		log.Printf("Warning: failed to disconnect from mongodb: %v", err)
	}

	return nil
}
